package com.cheetah.start.common.shoesImg;

import org.opencv.core.*;
import org.opencv.imgproc.Imgproc;
import org.opencv.imgcodecs.Imgcodecs;
import java.util.ArrayList;
import java.util.List;

public class UniversalBackgroundRemoval {

    /**
     * 自动检测并去除任何颜色的背景
     */
    public static Mat removeAnyBackground(Mat image) {
        // 方法1：尝试边缘检测+轮廓分析
        Mat result = removeByEdgesAndContours(image);

        // 方法2：如果方法1效果不好，尝试 GrabCut 算法
        if (isResultEmpty(result)) {
            result = removeByGrabCut(image);
        }

        // 方法3：如果还是不行，使用自适应阈值
        if (isResultEmpty(result)) {
            result = removeByAdaptiveThreshold(image);
        }

        return result;
    }

    /**
     * 基于边缘和轮廓的背景去除（最通用）
     */
    private static Mat removeByEdgesAndContours(Mat image) {
        Mat gray = new Mat();
        Mat blurred = new Mat();
        Mat edges = new Mat();

        // 转为灰度并降噪
        Imgproc.cvtColor(image, gray, Imgproc.COLOR_BGR2GRAY);
        Imgproc.GaussianBlur(gray, blurred, new Size(5, 5), 1.5);

        // 自适应Canny边缘检测
        double median = getMedianValue(blurred);
        int lowerThreshold = (int) Math.max(0, 0.7 * median);
        int upperThreshold = (int) Math.min(255, 1.3 * median);
        Imgproc.Canny(blurred, edges, lowerThreshold, upperThreshold);

        // 膨胀连接边缘
        Mat kernel = Imgproc.getStructuringElement(Imgproc.MORPH_ELLIPSE, new Size(3, 3));
        Imgproc.dilate(edges, edges, kernel, new Point(-1, -1), 2);

        // 查找轮廓
        List<MatOfPoint> contours = new ArrayList<>();
        Mat hierarchy = new Mat();
        Imgproc.findContours(edges, contours, hierarchy,
                Imgproc.RETR_EXTERNAL, Imgproc.CHAIN_APPROX_SIMPLE);

        // 创建掩码
        Mat mask = Mat.zeros(edges.size(), CvType.CV_8UC1);

        // 找到面积最大的几个轮廓
        List<MatOfPoint> largeContours = new ArrayList<>();
        for (MatOfPoint contour : contours) {
            double area = Imgproc.contourArea(contour);
            if (area > 1000) { // 根据图片大小调整
                largeContours.add(contour);
            }
        }

        // 填充轮廓创建掩码
        if (!largeContours.isEmpty()) {
            Imgproc.fillPoly(mask, largeContours, new Scalar(255));

            // 形态学操作平滑边界
            Mat morphKernel = Imgproc.getStructuringElement(Imgproc.MORPH_ELLIPSE, new Size(7, 7));
            Imgproc.morphologyEx(mask, mask, Imgproc.MORPH_CLOSE, morphKernel);
            Imgproc.morphologyEx(mask, mask, Imgproc.MORPH_OPEN, morphKernel);
        }

        // 应用掩码
        Mat result = new Mat();
        Core.bitwise_and(image, image, result, mask);

        return result;
    }

    /**
     * GrabCut算法 - 智能前景提取
     */
    private static Mat removeByGrabCut(Mat image) {
        Mat result = new Mat();
        Mat mask = new Mat();
        Mat bgdModel = new Mat();
        Mat fgdModel = new Mat();

        // 初始化矩形（假设鞋子在图片中央）
        Rect rect = new Rect(50, 50, image.cols() - 100, image.rows() - 100);

        // 初始化掩码
        mask.create(image.size(), CvType.CV_8UC1);
        mask.setTo(new Scalar(Imgproc.GC_PR_BGD));

        // 设置矩形内为可能的前景
        Mat rectMask = new Mat(mask, rect);
        rectMask.setTo(new Scalar(Imgproc.GC_PR_FGD));

        // 执行GrabCut
        Imgproc.grabCut(image, mask, rect, bgdModel, fgdModel, 3, Imgproc.GC_INIT_WITH_MASK);

        // 创建最终掩码
        Mat foregroundMask = new Mat();
        Core.compare(mask, new Scalar(Imgproc.GC_FGD), foregroundMask, Core.CMP_EQ);
        Core.compare(mask, new Scalar(Imgproc.GC_PR_FGD), mask, Core.CMP_EQ);
        Core.bitwise_or(foregroundMask, mask, foregroundMask);

        // 应用掩码
        image.copyTo(result, foregroundMask);

        return result;
    }

    /**
     * 自适应阈值方法
     */
    private static Mat removeByAdaptiveThreshold(Mat image) {
        Mat gray = new Mat();
        Mat binary = new Mat();

        Imgproc.cvtColor(image, gray, Imgproc.COLOR_BGR2GRAY);

        // 自适应阈值处理
        Imgproc.adaptiveThreshold(gray, binary, 255,
                Imgproc.ADAPTIVE_THRESH_GAUSSIAN_C,
                Imgproc.THRESH_BINARY, 11, 2);

        // 形态学操作
        Mat kernel = Imgproc.getStructuringElement(Imgproc.MORPH_ELLIPSE, new Size(5, 5));
        Imgproc.morphologyEx(binary, binary, Imgproc.MORPH_CLOSE, kernel);
        Imgproc.morphologyEx(binary, binary, Imgproc.MORPH_OPEN, kernel);

        Mat result = new Mat();
        Core.bitwise_and(image, image, result, binary);

        return result;
    }

    private static double getMedianValue(Mat mat) {
        Mat sorted = new Mat();
        Core.sort(mat.reshape(1, 1), sorted, Core.SORT_ASCENDING);
        return sorted.get(0, sorted.cols() / 2)[0];
    }

    private static boolean isResultEmpty(Mat result) {
        return Core.countNonZero(result.reshape(1, 1)) < 100;
    }
}
